# Modify Date: 2021/6/22
# Usage: output connected line include row and col
# Modify Date: 2021/8/13
# Usage: Connect gap from the same line
# Modify Date: 2021/11/12
# Usage: about test image, process segment map(remove noise and connect broken line)
# Modify Date: 2021/11/22
# Usage: add line connection-傅里叶变换->圆周方向累加->最大能量方向->扩展方向范围->傅里叶反变换
# Modify Date: 2021/11/24
# Usage: add line connection-最大方向范围缩减5->3,二值化图像处理
# Modify Date: 2021/12/2
# Usage: 计算FFT图圆周方向累加识别选取倾斜线周围像素再反变换
import os
import cv2
import math
import numpy as np
import matplotlib.pyplot as plt

def line_row_gen(img_path, out_path):
    print('当前图像:', img_path)
    img = cv2.imread(img_path)
    img_mask = np.ones_like(img) * 255
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    ''' 2021/12/2 add--傅里叶变换->正经圆周方向累加->最大能量方向->扩展方向范围->傅里叶反变换 '''
    f = np.fft.fft2(gray)
    fshift = np.fft.fftshift(f)
    magnitude_spectrum = np.log(np.abs(fshift))
    plt.subplot(131)
    plt.imshow(magnitude_spectrum, cmap='gray')
    plt.title('FFT')
    h, w = magnitude_spectrum.shape
    c_x, c_y = h//2-1, w//2-1
    angle = np.linspace(0, np.pi, 181)
    # l = np.sqrt(math.pow(h, 2)+math.pow(w, 2))
    r = np.linspace(-h/2, h/2, 257)
    point = []
    s_point = []
    start_sum, sum = 0, 0
    for a in angle:
        point.clear()
        for r_ in r:
            x_ = int(np.round(c_x+r_*math.cos(a)))
            y_ = int(np.round(c_y+r_*math.sin(a)))
            if x_ > w-1 or x_ < 0 or y_ > h-1 or y_ < 0:
                continue
            point.append((y_, x_))
        point = list(set(point))
        for point_ in point:
            sum += magnitude_spectrum[point_]
        # print(a, sum, start_sum)
        if sum > start_sum:
            start_sum = sum
            s_point = point
        sum = 0
    # print(s_point)
    s_e_point = []
    for s_point_ in s_point:
        for i in range(5):
            if s_point_[1] - i >= 0:
                s_e_point.append((s_point_[0], s_point_[1] - i))
            if s_point_[1] + i <= h-1:
                s_e_point.append((s_point_[0], s_point_[1] + i))
        s_e_point.append(s_point_)
    for x in range(fshift.shape[0]):
        for y in range(fshift.shape[1]):
            if (x, y) not in s_e_point:
                fshift[x, y] = 0j
                magnitude_spectrum[x, y] = 0
    plt.subplot(132)
    plt.imshow(magnitude_spectrum, cmap='gray')
    plt.title('max energy')
    iimg = np.fft.ifft2(fshift)
    magnitude_spectrum = np.abs(iimg)
    plt.subplot(133)
    plt.imshow(magnitude_spectrum, cmap='gray')
    plt.title('inverse FFT')
    plt.show()
    m_max = max(map(max, magnitude_spectrum))
    m_min = min(map(min, magnitude_spectrum))
    for x in range(magnitude_spectrum.shape[0]):
        for y in range(magnitude_spectrum.shape[1]):
            magnitude_spectrum[x, y] = np.round(255.0*(magnitude_spectrum[x, y]-m_min)/(m_max-m_min))
    magnitude_spectrum = magnitude_spectrum.astype(np.uint8)

    _, magnitude_spectrum = cv2.threshold(magnitude_spectrum, 127, 255, cv2.THRESH_BINARY|cv2.THRESH_OTSU)
    # cv2.namedWindow('binary image', cv2.WINDOW_NORMAL)
    # cv2.imshow('binary image', magnitude_spectrum)
    # cv2.waitKey(0)
    # cv2.destroyWindow('binary image')

    row_contours, _ = cv2.findContours(magnitude_spectrum, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    print(len(row_contours))
    result = cv2.drawContours(img_temp, row_contours, -1, 0, thickness=-1)

    # cv2.namedWindow(img_path, cv2.WINDOW_NORMAL)
    # cv2.imshow(img_path, result)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    '''检测轮廓'''
    # row_contours, _ = cv2.findContours(binary_erode, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    # sorted_row_contours = sorted(row_contours, key=lambda x: cv2.contourArea(x), reverse=True)
    # max_contour = sorted_row_contours[0]
    # rect = cv2.minAreaRect(max_contour)
    # angle = rect[2]
    # print(angle)
    # img = Image.fromarray(binary_raw)
    # if angle < 15:
    #     img_rot = img.rotate(angle)
    # elif 75 < angle < 100:
    #     img_rot = img.rotate(angle-90)
    # plt.figure()
    # plt.subplot(1, 2, 1)
    # plt.imshow(img)
    # plt.subplot(1, 2, 2)
    # plt.imshow(img_rot)
    # plt.show()

    # ret, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)

    '''检测图中直线'''
    # minLineLength = 100
    # maxLineGap = 100
    # lines = cv2.HoughLinesP( binary,1,np.pi/180,100,minLineLength=minLineLength,maxLineGap=maxLineGap )
    # print('lines:', len(lines))
    # try:
    #     for line in lines:
    #         for x1,y1,x2,y2 in line:
    #             # cv2.line( img,( x1,y1 ),( x2,y2 ),( 0,255,0 ),2 )
    #             cv2.line(img_temp, (x1, y1), (x2, y2), (0,), 1)
    # except:
    #     return img_temp, 0
    # cv2.imwrite( 'E:/image/myhoughlinesp.jpg',img )

    cv2.imwrite(out_path, result)

    return img_temp, 1

def line_col_gen(img_path, out_path):
    img = cv2.imread(img_path)
    img_temp = np.ones_like(img) * 255

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 5))
    # gray = cv2.morphologyEx(gray, cv2.MORPH_ERODE, kernel, iterations=10)

    # cv2.namedWindow(img_path, cv2.WINDOW_NORMAL)
    # cv2.imshow(img_path, gray)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 6))
    # gray = cv2.morphologyEx(gray, cv2.MORPH_DILATE, kernel, iterations=10)
    gray = cv2.morphologyEx(gray, cv2.MORPH_DILATE, kernel, iterations=2)

    # cv2.namedWindow(img_path, cv2.WINDOW_NORMAL)
    # cv2.imshow(img_path, gray)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    ret, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)

    cv2.imwrite(out_path, binary)

    return img_temp, 1

def angle(v1, v2):
  dx1 = v1[2] - v1[0]
  dy1 = v1[3] - v1[1]
  dx2 = v2[2] - v2[0]
  dy2 = v2[3] - v2[1]
  angle1 = math.atan2(dy1, dx1)
  angle1 = int(angle1 * 180/math.pi)
  angle2 = math.atan2(dy2, dx2)
  angle2 = int(angle2 * 180/math.pi)
  if angle1*angle2 >= 0:
    included_angle = abs(angle1-angle2)
  else:
    included_angle = abs(angle1) + abs(angle2)
    if included_angle > 180:
      included_angle = 360 - included_angle
  return included_angle

def line_col_gen_houghline(img_path, out_path):
    AB = [0,0,100,0]

    img = cv2.imread( img_path )
    img_temp = np.ones_like(img) *255

    gray = cv2.cvtColor( img,cv2.COLOR_BGR2GRAY )
    ret, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV)

    cv2.namedWindow(img_path, cv2.WINDOW_NORMAL)
    cv2.imshow(img_path, gray)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    edges = cv2.Canny( gray,50,150,apertureSize = 3 )
    cv2.namedWindow(img_path, cv2.WINDOW_NORMAL)
    cv2.imshow(img_path, edges)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    minLineLength = 100
    maxLineGap = 100
    lines = cv2.HoughLinesP( binary,1,np.pi/180,100,minLineLength=20,maxLineGap=50 )
    try:
        for line in lines:
            for x1,y1,x2,y2 in line:
                CD = [x1, y1, x2, y2]
                angle_cross = angle(AB, CD)
                if angle_cross<90+15 and angle_cross>90-15:
                    cv2.line( img,( x1,y1 ),( x2,y2 ),( 255,0,0 ),2 )
                    # cv2.line(img_temp, (x1, y1), (x2, y2), (0,), 2)
    except:
        return img_temp, 0
    # points = [(box[0], box[1]), (box[2],box[1]), (box[2], box[3]), (box[0], box[3])]
    # cv2.fillPoly(image,[np.array(points)],(255,0,0))
    cv2.namedWindow(img_path, cv2.WINDOW_NORMAL)
    cv2.imshow(img_path, img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    breakpoint()
    cv2.imwrite(out_path, img)
    return img_temp, 1

def tx_post(row_path, nrow_path, col_path, ncol_path, save_row_path, save_nrow_path, save_col_path, save_ncol_path):
    row_image, is_row_exist = line_row_gen(row_path, save_row_path)
    nrow_image, is_nrow_exist = line_row_gen(nrow_path, save_nrow_path)
    col_image, is_col_exist = line_col_gen(col_path, save_col_path)
    ncol_image, is_ncol_exist = line_col_gen(ncol_path, save_ncol_path)


if __name__ == '__main__':
    col_root = r'.\tx_infer_data\col'
    row_root = r'.\tx_infer_data\row'
    ncol_root = r'.\tx_infer_data\ncol'
    nrow_root = r'.\tx_infer_data\nrow'
    save_col_root = r'.\tx_process_data\col'
    save_row_root = r'.\tx_process_data\row'
    save_ncol_root = r'.\tx_process_data\ncol'
    save_nrow_root = r'.\tx_process_data\nrow'

    img_names = os.listdir(col_root)
    for img_name in img_names:
        col_path = os.path.join(col_root, img_name)
        ncol_path = os.path.join(ncol_root, img_name)
        row_path = os.path.join(row_root, img_name)
        nrow_path = os.path.join(nrow_root, img_name)
        save_col_path = os.path.join(save_col_root, img_name)
        save_row_path = os.path.join(save_row_root, img_name)
        save_ncol_path = os.path.join(save_ncol_root, img_name)
        save_nrow_path = os.path.join(save_nrow_root, img_name)
        tx_post(row_path, nrow_path, col_path, ncol_path, save_row_path, save_nrow_path, save_col_path, save_ncol_path)
